import cv2
import copy
from torchvision import transforms
import torch
import numpy as np
import torch.nn as nn
import torch.utils.data as Data
import torch.optim as optim
from device import DEVICE
from models.yolo import YoloBody
from transform_tools import TransformTools, create_mask_image

CUDA_LAUNCH_BLOCKING = 1
import torch.optim.lr_scheduler as lr_scheduler

tt = TransformTools()


def make_image_data(img_file, crop_region, resize_h, resize_w):
    img = cv2.imread(img_file, cv2.IMREAD_COLOR)
    src_img_tensor = tt.img2tensor(img)

    tgt_img_tensor = copy.deepcopy(src_img_tensor)
    # 使用转换对象调整图像尺寸
    src_img_tensor = tt.resize_transform(src_img_tensor)
    tgt_img_tensor = tt.resize_transform(tgt_img_tensor)

    crop_region = [
        int(crop_region[0] * (resize_h / tgt_img_tensor.size()[1])),
        int(crop_region[1] * (resize_h / tgt_img_tensor.size()[1])),
        int(crop_region[2] * (resize_w / tgt_img_tensor.size()[2])),
        int(crop_region[3] * (resize_w / tgt_img_tensor.size()[2]))
    ]

    resize_target_transform = transforms.Resize((resize_h, resize_w))
    tgt_img_tensor = resize_target_transform(tgt_img_tensor)

    # image2 = torch.zeros_like(tgt_img_tensor)
    # image2[:, crop_region[0]:crop_region[1], crop_region[2]:crop_region[3]] = 1
    # tgt_img_tensor = image2[0].unsqueeze(0)

    # y1 y2 x1 x2
    tgt_img_tensor = create_mask_image(tgt_img_tensor, crop_region)
    return src_img_tensor, tgt_img_tensor


resize_h, resize_w = 128, 96
s1, t1 = make_image_data("1.jpg", (200, 430, 500, 650), resize_h, resize_w)
s2, t2 = make_image_data("2.jpg", (200, 430, 500, 650), resize_h, resize_w)
s3, t3 = make_image_data("3.jpg", (200, 430, 500, 650), resize_h, resize_w)
s4, t4 = make_image_data("4.jpg", (200, 430, 500, 650), resize_h, resize_w)

# batch_src_img_tensor = torch.cat((s1.unsqueeze(0), s2.unsqueeze(0), s3.unsqueeze(0), s4.unsqueeze(0)), dim=0)
# batch_tgt_img_tensor = torch.cat((t2.unsqueeze(0), t3.unsqueeze(0), t4.unsqueeze(0), t1.unsqueeze(0)), dim=0)

batch_src_img_tensor = torch.cat((s3.unsqueeze(0), s3.unsqueeze(0), s3.unsqueeze(0), s3.unsqueeze(0)), dim=0)
batch_tgt_img_tensor = torch.cat((t1.unsqueeze(0), t2.unsqueeze(0), t3.unsqueeze(0), t4.unsqueeze(0)), dim=0)

# print(batch_src_img_tensor.shape,batch_tgt_img_tensor.shape)

model = YoloBody(4).to(DEVICE)


class MyDataSet(Data.Dataset):
    """自定义DataLoader"""

    def __init__(self, inputs, outputs):
        super(MyDataSet, self).__init__()
        self.inputs = inputs
        self.outputs = outputs

    def __len__(self):
        return self.inputs.shape[0]

    def __getitem__(self, idx):
        return self.inputs[idx], self.outputs[idx]


loader = Data.DataLoader(MyDataSet(batch_src_img_tensor, batch_tgt_img_tensor), 2, False)

for inputs, targets in loader:
    print(inputs.shape, targets.shape)

#
#     src_img = tt.tensor2img(inputs[1].to("cpu"))
#     src_img.show()
#     tgt_ime =   tt.tensor2img(targets[1].to("cpu"))
#     tgt_ime.show()
#
#     src_img = tt.tensor2img(inputs[2].to("cpu"))
#     src_img.show()
#     tgt_ime =   tt.tensor2img(targets[2].to("cpu"))
#     tgt_ime.show()
#
#     src_img = tt.tensor2img(inputs[3].to("cpu"))
#     src_img.show()
#     tgt_ime =   tt.tensor2img(targets[3].to("cpu"))
#     tgt_ime.show()


epochs = 3000
# loss_fn = torch.nn.CrossEntropyLoss()
# loss_fn =  torch.nn.BCELoss()
loss_fn = torch.nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=1e-1, momentum=0.99)
# optimizer = torch.optim.Adam(muti_model.parameters(), lr=0.01)
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.99)
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2, verbose=True,
                                           threshold=0.0001, threshold_mode='rel', cooldown=0, min_lr=0, eps=1e-08)

for epoch in range(epochs):
    for inputs, targets in loader:
        inputs, targets = inputs.to(DEVICE), targets.to(DEVICE)
        output = model(inputs)
        # print(output.shape,targets.shape)
        loss = loss_fn(output, targets)
        # print("Epoch {} {} {}".format(epoch + 1, loss, scheduler.get_last_lr()[0]))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    print("Epoch {} {} ".format(epoch + 1, loss))
    if (epoch + 1) % 50 == 0:
        scheduler.step(loss)

torch.save(model, "senet.pth")
